import datetime
import logging

from absl import app

from coin.base.datetime_util import to_datetime
from coin.feed.fastfeed.native_wrapper import run_from_feed_cache

from experimental.prophet.graph import graph
from experimental.prophet.graph.shape import Shape

from experimental.prophet.graph_runner.dataframe import run_from_dataframe
from experimental.prophet.graph_runner.fastfeed import (run_from_fastfeed_multiprocess,
                                                        FastfeedGraphAdapter)

from experimental.prophet.ops import globals
from experimental.prophet.ops.constant import constant
from experimental.prophet.ops.placeholder import placeholder
from experimental.prophet.ops.fastfeed import fastfeed_coin
from experimental.prophet.ops import timeseries
from experimental.prophet.ops import y_gen
from experimental.prophet.ops import elemwise_math
from experimental.prophet.ops.control_flow import cond
from experimental.prophet.ops.aggregator import scalar_value_aggregator
from experimental.prophet.ops.timer import timer
from experimental.prophet.ops.debug import printer

from experimental.prophet.craps.universe import universe_btc


def symbol_basic(feed,
                 nickname,
                 inverse=False,
                 price_multiplier=1.,
                 windows=None,
                 periodic_timer=None):
  windows = windows or ['1s', '2s', '5s', '10s', '30s', '1m', '5m', '10m', '30m', '60m']
  periodic_timer = periodic_timer or timer('1s')

  with graph.namespace(nickname):
    book = graph.create_node('fastfeed.BookBestLevel', [feed.book])
    trade = graph.create_node('fastfeed.Trade', [feed.trade])

    last_book_time = book.outputs[0]
    last_trade_time = trade.outputs[0]
    feed_ts = elemwise_math.max(elemwise_math.max(last_book_time, last_trade_time), periodic_timer)

    with graph.as_default_timestamp(feed_ts):
      ask0_p = (book.outputs[1] * price_multiplier).named_ref('b.ask0_p')
      ask0_q = book.outputs[2].named_ref('b.ask0_q')
      bid0_p = (book.outputs[3] * price_multiplier).named_ref('b.bid0_p')
      bid0_q = book.outputs[4].named_ref('b.bid0_q')
      spread = (ask0_p - bid0_p).named_ref('b.spread')
      mid_p = ((ask0_p + bid0_p) / 2.).named_ref('b.mid_p')
      mid_p_std = timeseries.time_moving_window(mid_p, windows).std()

      mid_p_ret = []
      y_mid_p_ret = []
      for w in windows:
        v = (mid_p - timeseries.time_shift(mid_p, w)).named_ref('mid_ret.' + w)
        mid_p_ret.append(v)
        y_mid_p_ret.append(y_gen.time_shift_y(v, w, timestamp=globals.global_timestamp()))

      book_vars = [
          ask0_p, ask0_q, bid0_p, bid0_q, spread, mid_p, *mid_p_std, *mid_p_ret, *y_mid_p_ret
      ]

      tp = (trade.outputs[1] * price_multiplier).named_ref('t.p')
      if not inverse:
        tq = trade.outputs[2].named_ref('t.q')
        tv = (tp * tq).named_ref('t.v')
      else:
        tq = (trade.outputs[2] / tp).named_ref('t.q')
        tv = trade.outputs[2].named_ref('t.v')
      tside = trade.outputs[3].named_ref('t.side')

      with graph.control_if(tside == 1):
        buy_q_sum = timeseries.time_moving_window(tq.named_ref('t.buy.q'), windows).sum()
        buy_v_sum = timeseries.time_moving_window(tv.named_ref('t.buy.v'), windows).sum()

      with graph.control_if(tside == 2):
        sell_q_sum = timeseries.time_moving_window(tq.named_ref('t.sell.q'), windows).sum()
        sell_v_sum = timeseries.time_moving_window(tv.named_ref('t.sell.v'), windows).sum()

      tq_sum = timeseries.time_moving_window(tq, windows).sum()
      tv_sum = timeseries.time_moving_window(tv, windows).sum()
      tp_std = timeseries.time_moving_window(tp, windows).std()
      trade_vars = [*tq_sum, *tv_sum, *tp_std, *buy_q_sum, *buy_v_sum, *sell_q_sum, *sell_v_sum]

  # TODO: last_trade_time > 5m, make it all None
  return book_vars + trade_vars


def model(date, machine=None):
  machine = machine or 'feed-01.ap-northeast-1.aws'
  ts = globals.global_timestamp()

  vars = []
  for symbol_info in universe_btc(date)[:10]:
    # if symbol_info.nickname not in ('bitmex-xbtusd', 'okex-btcusd-quarter', 'okex-btcusd-next_week'):
    #  continue

    feed = fastfeed_coin(symbol_info.product, symbol_info.sub_req, machine)
    vars += symbol_basic(feed,
                         symbol_info.nickname,
                         inverse=symbol_info.inverse,
                         price_multiplier=symbol_info.price_multiplier)

  with graph.control_if(ts % 10**9 == 0):
    aggregator = scalar_value_aggregator([ts, *vars])

  ts_1s = timer('1s')
  ts_30m = timer('30m')
  with graph.control_if(ts_30m == ts):
    printer('ts_30m:%t == global_ts:%t', ts_30m, ts, use_logger=True)

  return aggregator


def main(argv):
  '''
  g = graph.Graph()
  with graph.as_default_graph(g):
    model(datetime.date(2019, 4, 1))

  adapter = FastfeedGraphAdapter(g, 'feed-01.ap-northeast-1.aws')
  feed_sub_req = adapter.get_feed_sub_requests()

  run_from_feed_cache(
      feed_sub_req,
      adapter.on_book_reset,
      datetime.date(2019, 4, 1),
      datetime.date(2019, 4, 2),
      machine=adapter.default_machine)

  '''
  df = run_from_fastfeed_multiprocess(model,
                                      datetime.date(2019, 4, 1),
                                      datetime.date(2019, 4, 2),
                                      machine='feed-01.ap-northeast-1.aws',
                                      max_workers=None,
                                      use_run_cache=True,
                                      gen_dataframe_output=True,
                                      eval_updated_only=True,
                                      dump_pending_row=True,
                                      debug_executor=True,
                                      inject_date_to_graph_func=True)
  print(df)
  df.to_pickle('res-2.pkl')


if __name__ == '__main__':
  logging.basicConfig(
      level='DEBUG',
      format='%(levelname)8s %(asctime)s %(name)s %(filename)s:%(lineno)d] %(message)s')
  logging.getLogger('coin.exchange.util.feed_checker').setLevel(logging.WARNING)
  app.run(main)
